import torch
from tqdm import tqdm
from einops import rearrange

# 假设features和coords_grid已经定义好
features = torch.randn(20, 8107, 1024)  # 示例数据
coords_grid = torch.randn(8107, 2)    # 示例数据
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def batch_cal_sim(features, coords_grid):
    features = features.to(device)
    coords_grid = coords_grid.to(device)

    # 计算所有时间步的特征范数
    features_norm = features.norm(dim=2, keepdim=True)  # (t x (h x w)) x 1

    # 计算所有时间步的特征相似度矩阵
    # 这里会挤爆内存，需要改进
    t = features.shape[0]
    for source_t in tqdm(range(t), desc="source time"):
        source_features = features[source_t].to(device) # (h x w) x c
        affinity = torch.einsum("nc, tmc->tnm", source_features, features)  # t x n x m
        affinity = affinity / torch.clamp(features_norm * features_norm.transpose(1, 2), min=1e-08)  # t x n x m

        # 计算每个源特征的最佳匹配目标特征索引
        affinity_source_max = torch.argmax(affinity, dim=2)  # t x n

        # 计算每个目标特征的最佳匹配源特征索引
        affinity_target_max = torch.argmax(affinity, dim=1)  # t x m

        # 生成一个矩阵，其中每个元素是时间步数
        t = features.shape[0]
        time_steps = torch.arange(t).view(t, 1, 1).expand(-1, affinity_source_max.shape[1], -1).to(device)

        # 使用torch.gather来收集最佳匹配的索引
        affinity_target_max_gathered = torch.gather(affinity_target_max, 1, affinity_source_max.unsqueeze(2)).squeeze(2)  # t x n

        # 筛选出相互最近邻的源特征索引
        source_bb_indices = (time_steps == affinity_target_max_gathered).nonzero(as_tuple=True)  # (t, n) 中满足条件的索引

        # 根据筛选出的索引收集坐标和相似度
        best_buddies = {}
        for source_t, source_f in tqdm(zip(source_bb_indices[0], source_bb_indices[1]), desc="source time"):
            target_t = affinity_source_max[source_t, source_f]
            source_coords = coords_grid[source_f]
            target_coords = coords_grid[target_t]
            affinity_value = affinity[source_t, source_f, target_t]
            
            best_buddies[f'{source_t}_{target_t}'] = {
                "source_coords": source_coords,
                "target_coords": target_coords,
                "cos_sims": affinity_value
            }

    # 注意：上面的代码假设affinity_source_max和affinity_target_max的形状是(t, n)和(t, m)，
    # 并且每个时间步的源特征和目标特征数量相同（即n==m）。如果n和m不同，则需要进行适当的调整。
# batch_cal_sim(features, coords_grid)

def create_meshgrid(h, w, step=7, patch_size=14, return_hw=False):
        start_coord = patch_size//2
        x = torch.arange(start_coord, w, step=step, device=device).float()
        y = torch.arange(start_coord, h, step=step, device=device).float()
        yy, xx = torch.meshgrid(y, x)
        xx = xx.reshape(-1)
        yy = yy.reshape(-1)
        grid = torch.stack([xx, yy], dim=-1)
        if return_hw:
            return grid, len(y), len(x)
        return grid

def create_boxes(coords_grid, filt_idx, box_size=30):
    filt_coords = coords_grid[filt_idx]
    xmin = filt_coords[:, :, 0] - box_size # B x topk
    xmax = filt_coords[:, :, 0] + box_size # B x topk
    ymin = filt_coords[:, :, 1] - box_size # B x topk
    ymax = filt_coords[:, :, 1] + box_size # B x topk
    # concat to get boxes shaped B x topk x 4
    boxes = torch.cat([xmin[:, :, None], ymin[:, :, None], xmax[:, :, None], ymax[:, :, None]], dim=-1) 
    return boxes

def cal_sim(features, coords_grid):
    best_buddies = {}
    coords_grid = create_meshgrid(h=476, w=854, step=7).to(device)
    
    features = torch.randn(143, 8107, 1024) # T x C x H x W
    features = rearrange(features, 't c h w -> t (h w) c')
    torch.cuda.empty_cache()

    t = features.shape[0]
    for source_t in tqdm(range(t), desc="source time"):
        for target_t in tqdm(range(t), desc="target time"):
            if source_t == target_t:
                continue
            
            source_features = features[source_t].to(device) # (h x w) x c
            target_features = features[target_t].to(device) # (h x w) x c
            feature_range = torch.arange(source_features.shape[0]).to(device)

            affinity = torch.einsum("nc,mc->nm", source_features, target_features) # n x m
            affinity = affinity / torch.clamp(source_features.norm(dim=1)[:, None] * target_features.norm(dim=1)[None, ...], min=1e-08) # n x m
            affinity_source_max = torch.argmax(affinity, dim=1) # n
            affinity_target_max = torch.argmax(affinity, dim=0) # m
            
            # affinity_source_top3, source_top3_indices = torch.topk(affinity, 3, dim=1) # n x 3
            # affinity_target_top3, target_top3_indices = torch.topk(affinity, 3, dim=0) # m x 3
            
            # target_top3_indices_expanded = target_top3_indices.unsqueeze(-1)
            # target_top3_indices_expanded = target_top3_indices_expanded.expand(-1, -1, 4)
            # top4_boxes = torch.gather(boxes, 1, target_top3_indices_expanded)
            # top01iou = torch.diag(box_iou(top4_boxes[:, 0, :], top4_boxes[:, 1, :]))
            # top12iou = torch.diag(box_iou(top4_boxes[:, 1, :], top4_boxes[:, 2, :]))
            # topiou_r = torch.clamp( top01iou/ (top12iou + 1e-9), 0, 1)
    
            # /筛选出那些在源特征和目标特征之间形成 “相互最近邻”（mutual nearest neighbors）的源特征索引。
            source_bb_indices = feature_range == affinity_target_max[affinity_source_max]
            target_bb_indices = affinity_source_max[source_bb_indices]
            """
            在寻找最好匹配点的过程中，“相互最近邻” 的条件限制了匹配点的数量。对于每一个源特征，都要在目标特征中找到与之最相似的目标特征；
            同时，对于找到的这个目标特征，其最相似的源特征也必须是最初的那个源特征，这样才会被认定为一个最好匹配点。"""
            source_coords = coords_grid[source_bb_indices]
            target_coords = coords_grid[target_bb_indices]
            affinities = affinity[feature_range[source_bb_indices], target_bb_indices]
            
            best_buddies[f'{source_t}_{target_t}'] = {
                "source_coords": source_coords,
                "target_coords": target_coords,
                "cos_sims": affinities
            }